
import codecs
import math
import os
import platform
import sys
import cv2
import numpy as np
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *

xxx = 1.2

CURSOR_DEFAULT = Qt.ArrowCursor
CURSOR_POINT = Qt.PointingHandCursor
CURSOR_DRAW = Qt.CrossCursor
CURSOR_MOVE = Qt.ClosedHandCursor
CURSOR_GRAB = Qt.OpenHandCursor

def array2pix(image, ismask=False):
    if len(image.shape) == 2:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGRA)

    if image.shape[2] == 1:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGRA)

    if image.shape[2] == 3:
        #image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2BGRA)
    #show1 = cv2.resize(image, (640, 480))
    if ismask:
        #image[:,:,3] = image[:,:,2]
        pass

    show1 = image
    #QtImg = QImage(show1.data, show1.shape[1], show1.shape[0], QImage.Format_BGR888)
    QtImg = QImage(
        show1.data, show1.shape[1], show1.shape[0], QImage.Format_ARGB32)
    pixmap = QPixmap(QtImg)
    return pixmap

def pix2array(pixmap, ismask=False):
    channels_count = 4
    image = pixmap.toImage()
    w = pixmap.width()
    h = pixmap.height()
    s = image.bits().asstring(w * h * channels_count)
    arr = np.fromstring(s, dtype=np.uint8).reshape((h, w, channels_count))
    if ismask:
        arr = arr[:, :, 2]
    return arr

def paintCursor(k):
    if k < 2:
        k += 1
    pix = QPixmap(k+1, k+1)
    pix.fill(Qt.transparent)  # 用透明色填充
    painter = QPainter(pix)
    # mypen = QPen(QColor(0, 0, 0), 1, Qt.DashDotLine)  # 创建画笔
    mypen = QPen(QColor(255, 255, 0), 1)  # 创建画笔
    painter.setPen(mypen)  # 画板设置画笔
    k2 = int(k/2)
    painter.drawEllipse(QPointF(k2, k2), k2, k2)
    painter.end()
    cursor = QCursor(pix, k2, k2)
    return cursor

class BaseDrawer:
    def __init__(self, widget):
        self.widget = widget
        self.widget.drawer.append(self)
        self.widget.setMouseTracking(True)

    def update(self):
        self.widget.update()

    def save(self):
        pass

    def setCursor(self, a):
        self.widget.setCursor(a)

class ImageDrawer(BaseDrawer):
    def __init__(self, widget, image=None):
        super(ImageDrawer, self).__init__(widget)
        self.press = False
        self.point_list = []
        self.move_list = []
        self.epsilon = 11.0
        self.highlightShape = -1
        self.highlightVertex = (-1, -1, 0)
        self.k = 5
        # self.setCursor(paintCursor(self.k))  # 设置鼠标形状为十字形
        self._cursor = CURSOR_DEFAULT
        #self._painter = QPainter()
        self.buts = 0
        self.mods = 0
        self.shapes = []
        self.draw_mask = True
        self.pressInfo = None
        self.clickInfo = None
        self.cur_shape = -1
        self.nouse_ctrl = 1
        self.image = None
        self.pixmap = None
        self.offset = None
        self.set_img(image)
    
    def set_img(self, image):
        self.image = image
        if image is not None:
            self.pixmap = array2pix(image)
        if self.offset is None:
            self.init_trans()
        return 0
    
    
    def init_trans(self):
        if self.image is None:
            return 0
        self.scale = 1
        self.offset = [0, 0]
        imgw, imgh = self.pixmap.width(), self.pixmap.height()
        aa = self.widget.size()
        winw, winh = aa.width(), aa.height()
        self.scale = min(winh/imgh, winw/imgw)
        self.offset[0] = (winw-self.scale*imgw)/2
        self.offset[1] = (winh-self.scale*imgh)/2

    def get_transform(self):
        if self.image is None:
            return 0
        transform = QTransform()
        transform = transform.translate(self.offset[0], self.offset[1])
        transform = transform.scale(self.scale, self.scale)
        return transform

    def draw(self, event: QPaintEvent, p:QPainter):
        if self.image is None:
            return 0
        #p = self._painter
        #pix = QPixmap.fromImage(self.QtImg)
        transform = self.get_transform()
        p.setTransform(transform)
        p.drawPixmap(0, 0, self.pixmap)
        p.setTransform(QTransform())
        #shape = Shape()
        # self.drawext(a0,p)
        # p.end()

    def inv_trans_point(self, point):
        if self.image is None:
            return 0
        transform = self.get_transform()
        ttt = transform.inverted()[0]
        #aa = ttt.m11()
        pos = ttt.map(point)
        return pos

    def hitTest(self, point):
        if self.image is None:
            return 0
        #pos = self.inv_trans_point(point)
        cursor = CURSOR_DEFAULT

        transform = self.get_transform()
        ttt = transform.inverted()[0]
        epsilon = ttt.m11() * self.epsilon
        pos = ttt.map(point)

        highlightShape = -1
        highlightVertex = (-1, -1, 0)

        return highlightShape, highlightVertex, cursor

    # 鼠标左键按下，开始记录坐标
    def mousePressEvent(self, event):
        if self.image is None:
            return 0
        self.mods = int(event.modifiers())
        self.buts = event.button()
        point = QPointF(event.pos())
        self.press = True
        highlightShape, highlightVertex, cursor = self.hitTest(point)
        self.pressInfo = (highlightShape, highlightVertex, cursor)
        self.prevPoint = point
        if(event.button() == Qt.LeftButton):
            if Qt.ControlModifier == self.mods or self.nouse_ctrl:
                self.overrideCursor(CURSOR_MOVE)
                self.move_list = [point]
                return

            if self.cur_shape >= 0:
                pos = self.inv_trans_point(point)
                prevPos = self.inv_trans_point(self.prevPoint)
                dp = pos - prevPos
                shape = self.shapes[self.cur_shape]
                shape.moveVertexBy(2, dp)
                self.save()
                self.cur_shape = -1

        self.update()  # 更新，让画笔画画
        self.updateCursor()

    # 鼠标松开，停止记录
    def mouseReleaseEvent(self, event):
        if self.image is None:
            return 0
        self.mods = int(event.modifiers())
        self.buts = 0
        # if (self.buts == Qt.LeftButton):
        self.move_list = []
        self.press = False
        self.clickInfo = self.pressInfo
        self.pressInfo = None
        self.updateCursor()

    # 鼠标松开，停止记录
    def mouseReleaseEvent(self, event):
        if self.image is None:
            return 0
        self.mods = int(event.modifiers())
        self.buts = 0
        # if (self.buts == Qt.LeftButton):
        self.move_list = []
        self.press = False
        self.clickInfo = self.pressInfo
        self.pressInfo = None
        self.updateCursor()
        self.save()

    # 鼠标移动事件

    def mouseMoveEvent(self, event):
        if self.image is None:
            return 0
        point = QPointF(event.pos())  # 记录坐标
        self.pos = point
        #print(point.x(), point.y())
        mods = event.modifiers()
        if(event.buttons() == Qt.LeftButton):
            if int(mods) == Qt.ControlModifier or self.nouse_ctrl:
                if len(self.move_list) > 0:
                    prept = self.move_list[-1]
                    self.offset[0] += (point.x()-prept.x())  # /self.scale
                    self.offset[1] += (point.y()-prept.y())  # /self.scale
                    self.move_list.append(point)
                    self.update()  # 更新，让画笔画画
                    return

                return

        if self.cur_shape >= 0 and self.cur_shape < len(self.shapes):
            pos = self.inv_trans_point(point)
            prevPos = self.inv_trans_point(self.prevPoint)
            dp = pos - prevPos
            shape = self.shapes[self.cur_shape]
            self.prevPoint = point
            shape.moveVertexBy(2, dp)
            self.update()  # 更新，让画笔画画
            return

        if self.press:  # 只有鼠标被按下的时候才记录坐标

            if self.pressInfo is not None:
                # print(self.pressInfo)
                highlightShape, highlightVertex, cursor = self.pressInfo
                pos = self.inv_trans_point(point)
                prevPos = self.inv_trans_point(self.prevPoint)
                dp = pos - prevPos
                if(event.buttons() == Qt.RightButton):
                    if highlightVertex[0] >= 0:
                        i, index, type = highlightVertex
                        shape = self.shapes[i]
                        point = shape.points[index]
                        angle = getAngle(shape.center, pos, point)
                        shape.rotate(angle)

                if(event.buttons() == Qt.LeftButton):
                    if highlightVertex[0] >= 0:
                        i, index, type = highlightVertex
                        shape = self.shapes[i]
                        shape.moveVertexBy(index, dp)
                    elif highlightShape >= 0:
                        shape = self.shapes[highlightShape]
                        shape.moveBy(dp)

                self.save()
                self.prevPoint = point

            self.update()  # 更新，让画笔画画
        else:
            do_update = False
            highlightShape, highlightVertex, cursor = self.hitTest(point)

            if highlightShape != self.highlightShape:
                self.highlightShape = highlightShape
                do_update = True

            if highlightVertex != self.highlightVertex:
                self.highlightVertex = highlightVertex
                do_update = True

            if do_update:
                self.update()
            self.setCursor(cursor)
        return

    def updateCursor(self):
        if (self.buts and Qt.LeftButton):
            if self.mods == Qt.ControlModifier or self.nouse_ctrl:
                return self.overrideCursor(CURSOR_MOVE)

        if self.mods == Qt.ControlModifier:
            return self.overrideCursor(CURSOR_GRAB)

        return self.restoreCursor()

    def overrideCursor(self, cursor):
        self.restoreCursor()
        self._cursor = cursor
        QApplication.setOverrideCursor(cursor)

    def restoreCursor(self):
        QApplication.restoreOverrideCursor()

    def goto_image(self, x):
        pass

    def keyPressEvent(self, event:QKeyEvent):  # 键盘某个键被按下时调用
        self.mods = event.modifiers()
        key = event.key()
        ctrl = int(self.mods) == Qt.ControlModifier

        if key == Qt.Key_Escape:
            self.shapes.pop(self.cur_shape)
            self.cur_shape = -1
            self.update()
            return

        if key == Qt.Key_PageDown:
            self.goto_image(100)

        if key == Qt.Key_PageUp:
            self.goto_image(-100)

        if ctrl and key == Qt.Key_Space:
            self.goto_image(10)

        if ctrl and key == Qt.Key_B:
            self.goto_image(-10)

        if key == Qt.Key_Space:
            self.goto_image(1)

        if key == Qt.Key_B:
            self.goto_image(-1)

        if key == Qt.Key_D:
            if self.highlightShape >= 0:
                self.shapes.pop(self.highlightShape)
                self.save()
                self.update()

        if key == Qt.Key_F:
            self.draw_mask = not self.draw_mask
            self.update()

        if key == Qt.Key_Q or key == Qt.Key_W:
            if len(self.labels) > 0:
                if self.highlightShape >= 0:
                    t = self.shapes[self.highlightShape].label
                    i = self.labels.index(t) if t in self.labels else 0
                    i = max(i, 0)
                    i += 1 if Qt.Key_W == key else -1
                    if i >= len(self.labels):
                        i = 0
                    if i < 0:
                        i = len(self.labels)-1
                    self.shapes[self.highlightShape].label = self.labels[i]
                    self.pre_label = self.labels[i]
                    self.save()
                    self.update()

        if self.clickInfo is not None:
            if key == Qt.Key_D or key == Qt.Key_Delete:
                highlightShape, highlightVertex, cursor = self.clickInfo
                if highlightShape < 0:
                    i, index, type = highlightVertex
                    highlightShape = i

                if highlightShape >= 0 and len(self.shapes) > 0 and highlightShape < len(self.shapes):
                    self.shapes.pop(highlightShape)
                    self.update()

        if key == Qt.Key_R:
            point = self.pos
            self.prevPoint = point
            pos = self.inv_trans_point(point)
            self.addItem(pos, False)

        self.updateCursor()

    def changeEvent(self, e):
        if e.type() == QEvent.WindowStateChange:
            if self.isMinimized():
                print("窗口最小化")
            elif self.isMaximized():
                print("窗口最大化")
                self.update()
                self.init_trans()
            elif self.isFullScreen():
                print("全屏显示")
            elif self.isActiveWindow():
                print("活动窗口")

    def keyReleaseEvent(self, event:QKeyEvent):
        # if QKeyEvent.modifiers()==Qt.ControlModifier:
        self.mods = event.modifiers()
        self.updateCursor()
        self.save()

    def zoomRequest(self, delta, pos):
        units = xxx
        if delta < 0:
            units = 1/units

        if self.offset is None:
            self.init_trans()
            
        self.offset[0] = self.offset[0]*units + (1-units)*pos.x()
        self.offset[1] = self.offset[1]*units + (1-units)*pos.y()
        self.scale *= units

    def updateCursor1(self, angle):
        if (angle > 0):  # 如果滚轮往上滚
            if self.k > 200:
                return
            self.k = int(self.k*xxx + 1)
            self.k = min(self.k, 200)

        if (angle < 0):  # 如果滚轮往上滚
            if self.k < 2:
                return
            self.k = int(self.k/xxx)
            self.k = max(self.k, 2)

        k = int(self.k * self.scale)
        if k % 2 == 0:
            k += 1

        self.setCursor(paintCursor(k))  # 设置鼠标形状为十字形
        self.update()  # 更新，让画笔画画

    def wheelEvent(self, event):
        point = QPointF(event.pos())
        mods = event.modifiers()
        delta = event.angleDelta()
        angle = event.angleDelta().y()

        if Qt.ControlModifier == self.mods or self.nouse_ctrl:
            # with Ctrl/Command key
            # zoom
            self.zoomRequest(delta.y(), point)
            self.update()  # 更新，让画笔画画
            self.updateCursor1(0)
            return

        self.updateCursor1(angle)

class QDrawWidget(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.drawer = []

    def clear(self):
        self.drawer = []

    # 绘画，每次update的时候执行
    def paintEvent(self, event: QPaintEvent):
        if len(self.drawer)==0:
            return super(QDrawWidget, self).paintEvent(event)

        if len(self.drawer)>0:
            p = QPainter(self)  # 创建画板，画板附着在对象this上，也就是Widget窗口
            #p.begin(self)
            for d in self.drawer:
                d.draw(event, p)
        
            #p.end()
        return

    # 鼠标左键按下，开始记录坐标
    def mousePressEvent(self, event):
        for d in self.drawer:
            d.mousePressEvent(event)

    # 鼠标松开，停止记录
    def mouseReleaseEvent(self, event):
        for d in self.drawer:
            d.mouseReleaseEvent(event)

    # 鼠标移动事件

    def mouseMoveEvent(self, event):
        for d in self.drawer:
            d.mouseMoveEvent(event)
        return

    def keyPressEvent(self, event:QKeyEvent):  # 键盘某个键被按下时调用
        for d in self.drawer:
            d.keyPressEvent(event)

    def keyReleaseEvent(self, event:QKeyEvent):
        for d in self.drawer:
            d.keyReleaseEvent(event)

    def wheelEvent(self, event): #
        for d in self.drawer:
            d.wheelEvent(event)

def is_debug():
    gettrace = getattr(sys, 'gettrace', None)
    if gettrace is None:
        return False
    elif gettrace():
        return True
    else:
        return False

from common.utils import *

class UI_ImageView(QDrawWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.imgdrawer = ImageDrawer(self)
    
    def set_img(self, img):
        #self.clear()
        self.imgdrawer.set_img(img)
        self.update()
        return 0

if __name__ == '__main__':
    app = QApplication(sys.argv)
    if len(sys.argv) > 1:
        pa = sys.argv[1]
    fn = r'D:\code\git\ywlydd\cvtestdata\halcon\aaa\20200307162538584.png'
    #fn = r'D:\data\save\gree_nengchipian1\20230818\OK\20230818_150604_659\img_20230818_150604_814_NOREAD_1_OK_img.jpg'
    img = cv_imread(fn)
    btn = UI_ImageView()
    btn.set_img(img)
    #btn.set_img(None)
    #imgdrawer = ImageDrawer(btn, img)
    btn.resize(800, 600)
    btn.show()
    sys.exit(app.exec_())

